import torch
from torch import nn
from alias_multinomial import AliasMethod
import math
import torch.nn.functional as F

class NCEAverage(nn.Module):

    def __init__(self, inputSize, outputSize, K, T=0.07, momentum=0.5, use_softmax=True):
        '''
        inputSize：输入特征的维度。 bit
        outputSize：输出特征的维度，即类别的数量。
        K：每个正样本对应的负样本数量。
        T：温度参数，用于缩放相似度分数。
        momentum：动量参数，用于更新内存中的样本表示。
        use_softmax：布尔值，指示是否使用softmax函数来计算相似度分数。

        register_buffer 函数用于在 torch.nn.Module 中注册一个持久的缓冲区。缓冲区是模型的一部分，但不会被视为模型的参数，因此不会在调用 model.parameters() 时返回。缓冲区通常用于保存一些固定的值或状态，例如均值和方差。
        '''
        super(NCEAverage, self).__init__()
        self.nLem = outputSize #类别数量
        self.unigrams = torch.ones(self.nLem) #用于初始化采样器的权重，初始化为全 1（均匀分布）。
        self.multinomial = AliasMethod(self.unigrams) #使用 AliasMethod 创建负样本采样器。
        self.multinomial.cuda()
        self.K = K #保存负样本数量。
        self.use_softmax = use_softmax
        #创建一个包含超参数的缓冲区，可以使用self.params来访问这些超参数
        self.register_buffer('params', torch.tensor([K, T * math.sqrt(inputSize), -1, -1, momentum]))
        stdv = 1. / math.sqrt(inputSize / 3)
        rnd = torch.randn(outputSize, inputSize).mul_(2 * stdv).add_(-stdv)
        #创建一个包含内存的缓冲区，用于保存类别的表示。每个向量的大小为 [outputSize, inputSize]。初始化为均匀分布的随机值（标准差为 stdv）。使用 F.normalize 对每个向量进行归一化，使得每行的 L2 范数为 1。
        self.register_buffer('memory', F.normalize(rnd.sign(), dim=1))


    def update_memory(self, data):
        '''
        输入 data 是一组样本的特征表示。
        步骤:
            遍历输入特征，将它们累加。
            对累加结果按 L2 范数进行归一化。
            使用动量更新内存：先将内存清零（mul_(0)），再将归一化结果赋值给 memory。
        '''
        memory = 0
        for i in range(len(data)):
            memory += data[i]
        memory /= memory.norm(dim=1, keepdim=True)
        self.memory.mul_(0).add_(memory)
    
    def forward(self, l, ab, y, idx=None, epoch=None):
        # images_outputs, texts_outputs, torch.cat(idx* len(images)), epoch=epoch-args.warmup_epoch
        '''
        l: 模态 1 的特征，例如图像特征或某个特征编码器的输出。
        ab: 模态 2 的特征，例如文本特征或另一种特征编码器的输出。
        y: 样本的类别标签，表示输入的正样本对应的索引。
        idx: 用于指定负样本索引。如果未提供，将从 self.multinomial 采样器生成。
        epoch: 当前的训练轮次，用于动态调整 momentum 参数。
        '''
        K = int(self.params[0].item())  # 负样本数量 K
        T = self.params[1].item()  # 温度参数
        Z_l = self.params[2].item()  # 模态 l 的归一化常数
        Z_ab = self.params[3].item()  # 模态 ab 的归一化常数

        momentum = self.params[4].item() if (epoch is None) else (0 if epoch < 0 else self.params[4].item())
        batchSize = l.size(0)  # 当前 batch 的样本数
        outputSize = self.memory.size(0)  # 内存大小（类别数）
        inputSize = self.memory.size(1)  # 特征维度

        # score computation
        if idx is None:
        # 若未提供索引，则从负样本采样器中生成 K 个负样本的索引。
            idx = self.multinomial.draw(batchSize * (self.K + 1)).view(batchSize, -1)

            idx.select(1, 0).copy_(y.data)
        # sample
        if momentum <= 0:
            weight = (l + ab) / 2. # 如果动量为 0，则直接使用当前输入特征
            inx = torch.stack([torch.arange(batchSize)] * batchSize)
            inx = torch.cat([torch.arange(batchSize).view([-1, 1]), inx[torch.eye(batchSize) == 0].view([batchSize, -1])], dim=1).to(weight.device).view([-1])
            weight = weight[inx].view([batchSize, batchSize, -1])
        else:
            weight = torch.index_select(self.memory, 0, idx.view(-1)).detach().view(batchSize, K + 1, inputSize)

        weight = weight.sign_()# 对权重取符号，增强数值稳定性

        out_ab = torch.bmm(weight, ab.view(batchSize, inputSize, 1))# 模态 ab 与内存中权重的相似度
        # sample
        out_l = torch.bmm(weight, l.view(batchSize, inputSize, 1)) # 模态 l 与内存中权重的相似度
        if self.use_softmax:
            out_ab = torch.div(out_ab, T)
            out_l = torch.div(out_l, T)
            out_l = out_l.contiguous()
            out_ab = out_ab.contiguous()
        else:
            out_ab = torch.exp(torch.div(out_ab, T))
            out_l = torch.exp(torch.div(out_l, T))
            # set Z_0 if haven't been set yet,
            # Z_0 is used as a constant approximation of Z, to scale the probs
            if Z_l < 0:
                self.params[2] = out_l.mean() * outputSize
                Z_l = self.params[2].clone().detach().item()
                print("normalization constant Z_l is set to {:.1f}".format(Z_l))
            if Z_ab < 0:
                self.params[3] = out_ab.mean() * outputSize
                Z_ab = self.params[3].clone().detach().item()
                print("normalization constant Z_ab is set to {:.1f}".format(Z_ab))
            # compute out_l, out_ab
            out_l = torch.div(out_l, Z_l).contiguous()
            out_ab = torch.div(out_ab, Z_ab).contiguous()

        # # update memory
        with torch.no_grad():
            l = (l + ab) / 2.
            l.div_(l.norm(dim=1, keepdim=True))
            l_pos = torch.index_select(self.memory, 0, y.view(-1))
            l_pos.mul_(momentum)
            l_pos.add_(torch.mul(l, 1 - momentum))
            l_pos = l_pos.div_(l_pos.norm(dim=1, keepdim=True))
            self.memory.index_copy_(0, y, l_pos)
        #返回两个模态的相似度分布：out_l 和 out_ab
        return out_l, out_ab
